### yianni's question ###

from worker import Worker
import matplotlib
matplotlib.use('Agg')
from pylab import *
#from numpy import *

class Question2(Worker):
    """ Which genre earned the most money, based on worldwide gross? Which genre earned the least? (Use the Smaller Data Set) """
    def __init__(self, framework):
        Worker.__init__(self, framework)
        self.title = 'Question 2'

    def get_result(self):
        ds = self.framework.large_dataset
        dict = {} # dict to store {genre:avg movie gross}
        genres = [] # list of all possible genres

        # cycle through each record, list every possible genre
        for record in ds:
            if record[3] not in genres:
                if len(record) == 28:
                    if record[-1]:
                        genres.append(record[3])

        # for each genre, cycle through each record, if genre matches then add its gross/count+1
        for genre in genres:
            gross_total = 0
            count = 0
            for record in ds:
                if record[3] == genre:
                    if len(record) == 28:
                        if record[-1]:     
                            gross = int(record[-1])
                            gross_total += gross
                            count += 1
            # calc the average gross for a movie of each genre, store in dict
            dict[genre] = gross_total / count

        # plot data with matplotlib module
        bar(range(0, len(dict.keys())), dict.values())
        title('Average Movie Gross for each Genre')
        xlabel('Genre')
        ylabel('Average Movie Gross')
        #ind = numpy.arange(len(dict.keys()))    
        #xticks(ind+len(dict.keys())/2, dict.keys())
        #yticks(numpy.arange(0, max(len(dict.values())), max(dict.values())/10))
        savefig('q2.png')
    
        #fig = gcf() # get current figure
    #return ('This graph shows the average movie gross for each genre',fig)
